{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# https://github.com/facebookresearch/mae/blob/main/models_mae.py#L177\n",
    "\n",
    "    def random_masking(self, x, mask_ratio):\n",
    "        \"\"\"\n",
    "        Perform per-sample random masking by per-sample shuffling.\n",
    "        Per-sample shuffling is done by argsort random noise.\n",
    "        x: [N, L, D], sequence\n",
    "        \"\"\"\n",
    "        N, L, D = x.shape  # batch, length, dim\n",
    "        len_keep = int(L * (1 - mask_ratio))\n",
    "        \n",
    "        noise = torch.rand(N, L, device=x.device)  # noise in [0, 1]\n",
    "        \n",
    "        # sort noise for each sample\n",
    "        ids_shuffle = torch.argsort(noise, dim=1)  # ascend: small is keep, large is remove\n",
    "        ids_restore = torch.argsort(ids_shuffle, dim=1)\n",
    "\n",
    "        # keep the first subset\n",
    "        ids_keep = ids_shuffle[:, :len_keep]\n",
    "        x_masked = torch.gather(x, dim=1, index=ids_keep.unsqueeze(-1).repeat(1, 1, D))\n",
    "\n",
    "        # generate the binary mask: 0 is keep, 1 is remove\n",
    "        mask = torch.ones([N, L], device=x.device)\n",
    "        mask[:, :len_keep] = 0\n",
    "        # unshuffle to get the binary mask\n",
    "        mask = torch.gather(mask, dim=1, index=ids_restore)\n",
    "\n",
    "        return x_masked, mask, ids_restore"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "https://github.com/Project-MONAI/research-contributions/tree/main/SwinUNETR/Pretrain\n",
    "class CrossCorrelationLoss(torch.nn.Module):\n",
    "    def __init__(self, args):\n",
    "        super().__init__()\n",
    "        self.args = args\n",
    "\n",
    "    def forward(self, z_i, z_j):\n",
    "        # empirical cross-correlation matrix\n",
    "        b, n, c = z_i.size()\n",
    "        loss = torch.zeros(1).cuda()\n",
    "        for batch_index in range(b):\n",
    "            c = z_i[batch_index].T @ z_j[batch_index]\n",
    "\n",
    "            c.div_(n)\n",
    "\n",
    "            on_diag = torch.diagonal(c).add_(-1).pow_(2).sum()\n",
    "            off_diag = self.off_diagonal(c).pow_(2).sum()\n",
    "            loss += (on_diag + self.args.lambd * off_diag)\n",
    "        return loss / b\n",
    "\n",
    "    @staticmethod\n",
    "    def off_diagonal(x):\n",
    "        # return a flattened view of the off-diagonal elements of a square matrix\n",
    "        n, m = x.shape\n",
    "        assert n == m\n",
    "        return x.flatten()[:-1].view(n - 1, n + 1)[:, 1:].flatten()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# https://github.com/facebookresearch/moco/tree/main\n",
    "    def forward(self, im_q, im_k):\n",
    "        \"\"\"\n",
    "        Input:\n",
    "            im_q: a batch of query images\n",
    "            im_k: a batch of key images\n",
    "        Output:\n",
    "            logits, targets\n",
    "        \"\"\"\n",
    "\n",
    "        # compute query features\n",
    "        q = self.encoder_q(im_q)  # queries: NxC\n",
    "        q = nn.functional.normalize(q, dim=1)\n",
    "\n",
    "        # compute key features\n",
    "        with torch.no_grad():  # no gradient to keys\n",
    "            self._momentum_update_key_encoder()  # update the key encoder\n",
    "\n",
    "            # shuffle for making use of BN\n",
    "            im_k, idx_unshuffle = self._batch_shuffle_ddp(im_k)\n",
    "\n",
    "            k = self.encoder_k(im_k)  # keys: NxC\n",
    "            k = nn.functional.normalize(k, dim=1)\n",
    "\n",
    "            # undo shuffle\n",
    "            k = self._batch_unshuffle_ddp(k, idx_unshuffle)\n",
    "\n",
    "        # compute logits\n",
    "        # Einstein sum is more intuitive\n",
    "        # positive logits: Nx1\n",
    "        l_pos = torch.einsum(\"nc,nc->n\", [q, k]).unsqueeze(-1)\n",
    "        # negative logits: NxK\n",
    "        l_neg = torch.einsum(\"nc,ck->nk\", [q, self.queue.clone().detach()])\n",
    "\n",
    "        # logits: Nx(1+K)\n",
    "        logits = torch.cat([l_pos, l_neg], dim=1)\n",
    "\n",
    "        # apply temperature\n",
    "        logits /= self.T\n",
    "\n",
    "        # labels: positive key indicators\n",
    "        labels = torch.zeros(logits.shape[0], dtype=torch.long).cuda()\n",
    "\n",
    "        # dequeue and enqueue\n",
    "        self._dequeue_and_enqueue(k)\n",
    "\n",
    "        return logits, labels\n",
    "    \n",
    "        criterion = nn.CrossEntropyLoss().cuda(args.gpu)\n",
    "                loss = criterion(output, target)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def info_nce_loss(self, features):\n",
    "\n",
    "    labels = torch.cat([torch.arange(self.args.batch_size) for i in range(self.args.n_views)], dim=0)\n",
    "    labels = (labels.unsqueeze(0) == labels.unsqueeze(1)).float()\n",
    "    labels = labels.to(self.args.device)\n",
    "\n",
    "    features = F.normalize(features, dim=1)\n",
    "\n",
    "    similarity_matrix = torch.matmul(features, features.T)\n",
    "    # assert similarity_matrix.shape == (\n",
    "    #     self.args.n_views * self.args.batch_size, self.args.n_views * self.args.batch_size)\n",
    "    # assert similarity_matrix.shape == labels.shape\n",
    "\n",
    "    # discard the main diagonal from both: labels and similarities matrix\n",
    "    mask = torch.eye(labels.shape[0], dtype=torch.bool).to(self.args.device)\n",
    "    labels = labels[~mask].view(labels.shape[0], -1)\n",
    "    similarity_matrix = similarity_matrix[~mask].view(similarity_matrix.shape[0], -1)\n",
    "    # assert similarity_matrix.shape == labels.shape\n",
    "\n",
    "    # select and combine multiple positives\n",
    "    positives = similarity_matrix[labels.bool()].view(labels.shape[0], -1)\n",
    "\n",
    "    # select only the negatives the negatives\n",
    "    negatives = similarity_matrix[~labels.bool()].view(similarity_matrix.shape[0], -1)\n",
    "\n",
    "    logits = torch.cat([positives, negatives], dim=1)\n",
    "    labels = torch.zeros(logits.shape[0], dtype=torch.long).to(self.args.device)\n",
    "\n",
    "    logits = logits / self.args.temperature\n",
    "    return logits, labels"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
